#include	"wynn.h"
using	namespace	wynn;

struct	Corpus {
    uint16_t	*ptr;
    uint64_t	len;
    int	fd;
    struct	stat	sb;
};

__global__	void	_emb(uint32_t	R,	uint32_t	C,	uint16_t	*inp,	half	*out) {
    uint32_t	id=blockIdx.x*blockDim.x+threadIdx.x,	r=id%R,	c=id/R;
    out[id]=__float2half_rn((wyhash64(inp[c],r)&1)*2-1.0f);
}

__global__	void	dlossf(uint32_t	C,	uint32_t	O,	half	*a,	uint16_t	*x,	float	*y,	uint16_t	*table) {
    uint32_t	i=blockIdx.x*blockDim.x+threadIdx.x;
    y[i]=-logf(fmaxf(__half2float(a[i*O+table[x[i+1]]]),FLT_MIN));
}

__global__	void	_zscore(uint32_t	C,float	*y) {
    double	sx=0,	sxx=0,	sz=0;
    for(uint32_t	i=0;	i<C;	i++) {
        float	x=y[C+i]-y[i];
        sx+=x;
        sxx+=x*x;
        sz+=y[i];
    }

    sx/=C;
    sxx=sqrt(sxx/C-sx*sx);
    y[0]=sz;
    y[1]=sqrtf(C)*sx/sxx;
}

template<uint32_t	C,	uint32_t	E,	uint32_t	D,	uint32_t	H,	uint32_t	S,	uint32_t	O>
struct	HomoSapiens {
  private:
    tensor<float,C*2>	ret;
    tensor<uint16_t,C+1>	data;
    tensor<uint16_t,O>	code;
    tensor<uint16_t,65536>	table;
    tensor<half,E*C>	n0g;
  public:
    tensor<half,E*C>	emb;
    linear<E,O,C>	ou;
    tensor<half,O*C>	out;

    void	load_voca(const	char	*F) {
        FILE	*f=fopen(F,"rt");
        for(uint32_t	i=0;	i<65536;	i++)	table.p[i]=0;
        for(uint32_t	i=0;	i<O;	i++) {
            fscanf(f,"%hu",code.p+i);
            table.p[code.p[i]]=i;
        }
        fclose(f);
    }

    bool	save(const	char	*F) {
        FILE	*f=fopen(F,"wb");
        if(f==NULL)	return	false;
        uint32_t	x;
        x=C;
        fwrite(&x,4,1,f);
        x=E;
        fwrite(&x,4,1,f);
        x=D;
        fwrite(&x,4,1,f);
        x=H;
        fwrite(&x,4,1,f);
        x=S;
        fwrite(&x,4,1,f);
        x=O;
        fwrite(&x,4,1,f);
        fwrite(code.p,O*2,1,f);
        ou.save(f);
        fclose(f);
        return	true;
    }

    bool	load(const	char	*F) {
        FILE	*f=fopen(F,"rb");
        if(f==NULL)	return	false;
        uint32_t	x;
        if(fread(&x,4,1,f)!=1||x!=C)	fprintf(stderr,"C=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=E)	fprintf(stderr,"E=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=D)	fprintf(stderr,"D=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=H)	fprintf(stderr,"H=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=S)	fprintf(stderr,"S=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=O)	fprintf(stderr,"O=%u\n",x);
        if(fread(code.p,O*2,1,f)!=1)	return	false;
        ou.load(f);
        fclose(f);
        return	true;
    }

    uint32_t	size(void) {
        return	ou.size();
    }

    void	forward(void) {
        uint32_t	w=warp_size(1,E*C);
        _emb<<<E*C/w,w>>>(E,C,data.p,emb.p);
        ou.fw(emb,out,0);
        _softmaxf<<<C/16,16>>>(O,out.p);
        w=warp_size(1,C);
        dlossf<<<C/w,w>>>(C,O,out.p,data.p,ret.p+stage*C,table.p);
    }

    float	train(vector<Corpus>	&corpus) {
        Corpus	*c=&corpus[wyrand(&prng)%corpus.size()];
        cudaMemcpy(data.p,c->ptr+(wyrand(&prng)%(c->len-C)),2*(C+1),cudaMemcpyHostToDevice);
        stage=0;
        forward();
        stage=1;
        forward();
        _zscore<<<1,1>>>(C,ret.p);
        cudaDeviceSynchronize();
        gradient=ret.p[1];
        ou.bk();
        cudaDeviceSynchronize();
        return	ret.p[0];
    }
};

#include	"config"

void	document(void) {
    cerr<<"usage:	training [options] input.txt\n";
    cerr<<"\t-i:	input model=NULL\n";
    cerr<<"\t-o:	output model=model\n";
    cerr<<"\t-s:	trained sample=0\n";
    cerr<<"\t-b:	benchmark only=off\n";
    exit(0);
}


int	main(int	ac,	char	**av) {
    HomoSapiens<context,embed,depth,headn,heads,voca>	model;
    string	in,out="model";
    int	opt,bench=0;
    uint64_t	training=0;
    while((opt=getopt(ac,	av,	"i:o:s:b"))>=0) {
        switch(opt) {
        case	'i':
            in=optarg;
            break;
        case	'o':
            out=optarg;
            break;
        case	's': {
            training=atoi(optarg);
            training<<=20;
        }
        break;
        case	'b':
            bench=1;
            prng=0;
            break;
        default:
            document();
        }
    }
    if(ac<optind+1) {
        document();
        return	0;
    }
    vector<Corpus>	vc(ac-optind);
    uint64_t	corpus_size=0;
    for(int	i=optind;	i<ac;	i++) {
        vc[i-optind].fd=open(av[i],	O_RDONLY);
        fstat(vc[i-optind].fd,	&vc[i-optind].sb);
        vc[i-optind].ptr=(uint16_t*)mmap(NULL,	vc[i-optind].sb.st_size,	PROT_READ,	MAP_SHARED,	vc[i-optind].fd,	0);
        corpus_size+=(vc[i-optind].len=vc[i-optind].sb.st_size/2);
        cerr<<av[i]<<'\t'<<vc[i-optind].len<<'\n';
    }
    cerr.precision(4);
    cerr.setf(ios::fixed);
    double	loss0=FLT_MAX/2,	loss;
    timeval	beg,	end;
    cerr<<"para\t"<<model.size()<<'\n';
    model.load_voca("voca.txt");
    if(in.size())	model.load(in.c_str());
    for(;;) {
        loss=0;
        gettimeofday(&beg,NULL);
        for(uint32_t	i=0;	i<fullbatch;	i++) {
            etalr=1.0f/sqrtf(model.size()*sqrtf(model.size()));
            training+=context;
            loss+=model.train(vc);
        }
        loss/=context*fullbatch;
        if(!bench) {
            if(loss<loss0+0.02)	model.save(out.c_str());
            else	break;
        }
        loss0=loss;
        gettimeofday(&end,NULL);
        double	t=(end.tv_sec-beg.tv_sec+1e-6*(end.tv_usec-beg.tv_usec));
        cerr<<(training>>20)<<'\t'<<loss<<'\t'<<t<<'\n';
    }
    for(int	i=optind;	i<ac;	i++) {
        munmap(vc[i-optind].ptr,vc[i-optind].sb.st_size);
        close(vc[i-optind].fd);
    }
    return	0;
}
